"""
模型的封装
"""
import fasttext

import config


class Classify:
    def __init__(self):
        # 加载模型
        self.model = fasttext.load_model(config.classify_model_path)
        self.model_by_word = fasttext.load_model(config.classify_model_by_word_path)

    def predict(self, sentence):
        """
        预测
        :param sentence: 字典{"cut": str, "cut_by_word": str}
        :return: (label, accuracy)
        """
        ret = self.model.predict(sentence["cut"])
        ret_by_word = self.model_by_word.predict(sentence["cut_by_word"])

        for label, acc, label_by_word, acc_by_word in zip(*ret, *ret_by_word):
            # 两个类别的情况
            if label == "__label__chat":
                label = "__label__qa"
                acc = 1 - acc

            if label_by_word == "__label__chat":
                label_by_word = "__label__qa"
                acc_by_word = 1 - acc_by_word

            # 判断准确率
            if acc > 0.95 or acc_by_word > 0.95:
                return "qa", max(acc, acc_by_word)
            else:
                return "chat", 1 - min(acc, acc_by_word)

            # # TODO 假设有三个类别
            # if label == label_by_word:
            #     if acc > 0.95 or acc_by_word > 0.95:
            #         return label, max(acc, acc_by_word)
            #     else:
            #         return None, 0  # 无法确认
            #
            # else:
            #     if acc_by_word > 0.99:
            #         return label_by_word, acc_by_word
            #     elif acc > 0.98:
            #         return label, acc
            #     else:
            #         return None, 0



